# python D:/pythonWork/KJ/ModelTrain.py D:/pythonWork/KJ/ 250
# -*- coding: UTF-8 -*-
from sklearn.neural_network import MLPRegressor
import joblib
import pandas as pd
import sys
import os

'''
训练多层神经感知网络
'''
def trainMlp(modelDir, splitIndex):
    # 解析文件，获取特征和结果
    train = pd.read_csv(f'{modelDir}/train.csv')
    with open(f'{modelDir}/labels.txt', 'w') as f:
        f.write(",".join(train.columns.tolist()))
        f.write("\n"+splitIndex)
    X = train.iloc[:, :int(splitIndex)]
    Y = train.iloc[:, int(splitIndex):]
    modelPath = f'{modelDir}/model.model'
    if os.path.exists(modelPath):
        mlp = joblib.load(modelPath)
        mlp.partial_fit(X, Y)
    else:
        mlp = MLPRegressor(hidden_layer_sizes=(128, 64),
                           activation='relu',
                           solver='adam',
                           max_iter=500,
                           random_state=42)
        mlp.fit(X, Y)
    # 生成模型文件
    joblib.dump(mlp, f'{modelDir}/model.model')


if __name__ == '__main__':
    # 获取参数
    modelDir = sys.argv[1]
    splitIndex = sys.argv[2]
#     modelDir = "D:\\work\\com.tixi.dky.kjtx.boot.be\\..\\upload\\dkyKjtx\\sceneModel"
    splitIndex = "30"
    trainMlp(modelDir, splitIndex)
